# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
import json
from collections import OrderedDict
from urllib import parse
from typing import Union, List

from tornado import gen, httpclient

from coin.base.param_util import to_list
import coin.flow.topic_map as topic_map
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.bitflyer_v1.config import group_symbol
from coin.exchange.bitflyer_v1.kr_rest.product import BitflyerProduct


def _generate_negotiate_url(account_id, token, symbol):
  data = OrderedDict()
  data['clientProtocol'] = 1.5
  data['account_id'] = account_id
  data['token'] = token
  data['products'] = '%s,heartbeat' % symbol
  data['connectionData'] = [{"name": "bfexhub"}]
  return 'https://signal.bitflyer.com/signalr/negotiate?%s' % parse.urlencode(data)


def _generate_connect_url(account_id, token, symbol, connection_token):
  data = OrderedDict()
  data['transport'] = 'webSockets'
  data['clientProtocol'] = 1.5
  data['account_id'] = account_id
  data['token'] = token
  data['products'] = '%s,heartbeat' % symbol
  data['connectionToken'] = connection_token
  data['connectionData'] = [{"name": "bfexhub"}]
  return 'wss://signal.bitflyer.com/signalr/connect?%s' % parse.urlencode(data)


def _generate_reconnect_url(account_id,
                            token,
                            symbol,
                            connection_token,
                            group_token,
                            last_message_id):
  data = OrderedDict()
  data['transport'] = 'webSockets'
  data['groupsToken'] = group_token
  data['messageId'] = last_message_id
  data['clientProtocol'] = 1.5
  data['account_id'] = account_id
  data['token'] = token
  data['products'] = '%s,heartbeat' % symbol
  data['connectionToken'] = connection_token
  data['connectionData'] = [{"name": "bfexhub"}]
  return 'wss://signal.bitflyer.com/signalr/reconnect?%s' % parse.urlencode(data)


def _generate_start_url(account_id, token, symbol, connection_token):
  data = OrderedDict()
  data['transport'] = 'webSockets'
  data['clientProtocol'] = 1.5
  data['account_id'] = account_id
  data['token'] = token
  data['products'] = '%s,heartbeat' % symbol
  data['connectionToken'] = connection_token
  data['connectionData'] = [{"name": "bfexhub"}]
  return 'https://signal.bitflyer.com/signalr/start?%s' % parse.urlencode(data)


class BitflyerSubscriber(WebsocketJsonSubscriberBase):
  BookFeedDisconnectWindowInNs = 2 * 60 * 1e9  # 2 mins

  def __init__(
      self,
      queue,
      group_id,
      worker_id,
      symbol,
      account_id='DEMO2',  # https://lightning.bitflyer.jp/home/demo?lang=en
      token='',  # empty string for DEMO2
      close_callback=None,
      logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    if isinstance(symbol, list):
      assert len(symbol) == 1, symbol
      symbol = symbol[0]
    assert isinstance(symbol, str), symbol
    self._symbol = symbol
    self._account_id = account_id
    self._token = token
    self._connection_token = None
    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)
    self._group_token = None
    self._last_message_id = None

  # Override base.
  # Bitflyer connection will be reset every 110 seconds, hide the reconnection event by ourself.
  # Only notify outside if reconnection related necessary info is missing.
  def _notify_close(self):
    if self._group_token and self._last_message_id and self._connection_token:
      if self._state != self.State.RUNNING:
        return
      self._state = self.State.CLOSED
      self._publish_topic_event('STOP', reason='CONNECTION_CLOSED')
      self._logger.info('Connection closed')
      if self._ws is not None:
        self._ws.close()

      self._register_url(
          _generate_reconnect_url(self._account_id,
                                  self._token,
                                  self._symbol,
                                  self._connection_token,
                                  self._group_token,
                                  self._last_message_id))
      self.start(self._ioloop)
    else:
      super()._notify_close()

  async def _negotiate(self):
    url = _generate_negotiate_url(self._account_id, self._token, self._symbol)
    req = httpclient.HTTPRequest(url=url, method='GET')
    res = await httpclient.AsyncHTTPClient().fetch(req)
    return json.loads(res.body)

  @gen.coroutine
  def _run_loop(self):
    negotiate_reuslt = yield self._negotiate()
    self._connection_token = negotiate_reuslt['ConnectionToken']
    self._register_url(
        _generate_connect_url(self._account_id, self._token, self._symbol, self._connection_token))
    super()._run_loop()

  async def _start(self):
    req = httpclient.HTTPRequest(url=_generate_start_url(self._account_id,
                                                         self._token,
                                                         self._symbol,
                                                         self._connection_token),
                                 method='GET')
    response = await httpclient.AsyncHTTPClient().fetch(req)

  @gen.coroutine
  def on_open(self):
    self._ioloop.add_timeout(datetime.timedelta(seconds=1), self._start)
    super().on_open()

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_ticker_board = topic_map.search_topic_id('bitflyer_ws_ticker_board_%s%s' %
                                                   (group_id, worker_id))[0]
    self._publisher = {
        'ticker_board': queue.get_writer(topic_id=topic_ticker_board[1]),
    }
    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: dict):
    if self.stopped:
      return
    if 'C' in message:
      self._last_message_id = message['C']
    elif message:
      logging.debug('bitflyer message: \n%s' % json.dumps(message, indent=2))
    if 'G' in message:
      self._group_token = message['G']
    self._publisher['ticker_board'].write(timestamp, message)


class BitflyerMultiSubscriber(object):
  BookFeedDisconnectWindowInNs = 2 * 60 * 1e9  # 2 mins

  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbols: Union[List[str], str],
               close_callback=None,
               logger=None):
    self._subscribers = []
    symbols = to_list(symbols)
    for symbol in symbols:
      if group_id == "TBD":
        matched_group_id = None
        for k, v in group_symbol.items():
          if symbol in v:
            matched_group_id = k
        assert matched_group_id is not None, (symbol, group_id)
      self._subscribers.append(
          BitflyerSubscriber(queue,
                             group_id,
                             worker_id,
                             symbol,
                             close_callback=close_callback,
                             logger=logger))

  def start(self, *args, **kwargs):
    for subscriber in self._subscribers:
      subscriber.start(*args, **kwargs)

  def stop(self, *args, **kwargs):
    for subscriber in self._subscribers:
      subscriber.stop(*args, **kwargs)
